Skip to content

[Unity][PyTorch] Disable gradient during dynamo subgraph capture to save RAM#14602

Merged
tqchen merged 1 commit intoapache:unityfrom
masahi:dynamo-no-grad
Apr 12, 2023
Merged

[Unity][PyTorch] Disable gradient during dynamo subgraph capture to save RAM#14602
tqchen merged 1 commit intoapache:unityfrom
masahi:dynamo-no-grad

Conversation

@masahi
Copy link
Copy Markdown
Member

@masahi masahi commented Apr 12, 2023

I got an out-of-memory error on my 16G linux system when trying to trace VAE in web-stable-diffusion https://github.com/mlc-ai/web-stable-diffusion/blob/main/build.py#L82 (this is not an issue on Mac). This simple change let me complete the build successfully.

Using the time -v command to measure the peak RAM usage of a script that does VAE tracing, I got

Before: Maximum resident set size (kbytes): 12288884
After (disable grad): Maximum resident set size (kbytes): 7179364

So this is a good reduction of 5GB.

@MasterJH5574 @spectrometerHBH @junrushao

@tvm-bot
Copy link
Copy Markdown
Collaborator

tvm-bot commented Apr 12, 2023

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

Generated by tvm-bot

@tqchen tqchen merged commit 88f5b8f into apache:unity Apr 12, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants